import os
from unstructured.partition.auto import partition
from unstructured.documents.elements import Title, NarrativeText
from typing import List

def process_document(file_path: str):
    """处理单个文档并提取内容"""
    try:
        print(f"\n{'='*60}")
        print(f"处理文件: {file_path}")
        print(f"{'='*60}")

        # 使用unstructured解析文档
        elements = partition(filename=file_path)

        # 提取不同类型的内容
        titles = []
        narratives = []
        other_content = []

        for element in elements:
            if isinstance(element, Title):
                titles.append(element.text)
            elif isinstance(element, NarrativeText):
                narratives.append(element.text)
            else:
                other_content.append(element.text)

        # 打印结果
        print("标题:")
        for i, title in enumerate(titles[:3], 1):  # 只显示前3个标题
            print(f"  {i}. {title}")
        if len(titles) > 3:
            print(f"  ... 还有 {len(titles) - 3} 个标题")

        print("\n正文内容:")
        for i, narrative in enumerate(narratives[:2], 1):  # 只显示前2段正文
            print(f"  {i}. {narrative[:200]}..." if len(narrative) > 200 else f"  {i}. {narrative}")
        if len(narratives) > 2:
            print(f"  ... 还有 {len(narratives) - 2} 段内容")

        print(f"\n总元素数量: {len(elements)}")
        print(f"标题数量: {len(titles)}")
        print(f"正文段落数量: {len(narratives)}")

    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {e}")

def process_html_from_path(filePath: str):
    """处理网页内容"""
    try:
        print(f"\n{'='*60}")
        print(f"处理网页: {filePath}")
        print(f"{'='*60}")

        # 解析网页
        elements = partition(filename=filePath)

        print("网页内容摘要:")
        for i, element in enumerate(elements[:5], 1):  # 只显示前5个元素
            element_type = type(element).__name__
            content_preview = element.text[:100] + "..." if len(element.text) > 100 else element.text
            print(f"  {i}. [{element_type}] {content_preview}")

        print(f"\n总元素数量: {len(elements)}")

    except Exception as e:
        print(f"处理网页 {filePath} 时出错: {e}")

def save_extracted_content(elements: List, output_file: str):
    """保存提取的内容到文件"""
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("提取的文档内容:\n")
        f.write("="*50 + "\n\n")

        for i, element in enumerate(elements, 1):
            element_type = type(element).__name__
            f.write(f"{i}. [{element_type}]\n")
            f.write(f"{element.text}\n")
            f.write("-"*50 + "\n\n")

def main():
    """主函数"""

    # 示例：你可以创建一些测试文件或使用现有文件
    # 这里只是示例，实际使用时请替换为真实文件路径

    print("Unstructured 文档解析示例")
    print("="*60)

    # 1. 处理网页
    print("\n1. 网页解析示例:")
    process_html_from_path("D:/tmp/36.云原生之SpringCloud+k8s实践-CSDN博客.html")

    # 2. 处理本地文件（需要实际文件路径）
    print("\n2. 本地文件解析示例:")

    #Markdown文件
    if (os.path.exists("D:/ideaSpace/rag-in-action-master/90-文档-Data/黑悟空/黑悟空版本介绍.md")):
        process_document("D:/ideaSpace/rag-in-action-master/90-文档-Data/黑悟空/黑悟空版本介绍.md")

    # PDF 文件
    if os.path.exists("D:/document/个人社保缴费清单.pdf"):
        process_document("D:/document/个人社保缴费清单.pdf")

    # Word 文档
    if os.path.exists("D:/document/胡龙福-430421198901257795.docx"):
        process_document("D:/document/胡龙福-430421198901257795.docx")

    # TXT 文件
    if os.path.exists("D:/document/新文件 1.txt"):
        process_document("D:/document/新文件 1.txt")

    # Excel 文件
    if os.path.exists("D:/document/工作文档/本地化清单V3-（0720日） - 韩玉飞交接.xlsx"):
        process_document("D:/document/工作文档/本地化清单V3-（0720日） - 韩玉飞交接.xlsx")

if __name__ == "__main__":
    main()